#include "loadidentify.hh"
#include <math.h>
#include <memory.h>
#include <iostream>
#include <stdio.h>
#include <fstream>
#include <eigen3/Eigen/Dense>
#include <vector>

#define PI 3.141592653
#define GRAVITY_CONST 9.81
#define MODIFY_FACTOR 1.2

// 粗略判断前后轨迹是否一致
#define TRAJ_DIFF_NUM 10

// 跳过头尾因滤波而畸变的数据
#define TRAJ_MULT_NUM 30

using namespace std;
using namespace Eigen;

static bool fileExist(const std::string& name)
{
    std::ifstream ifs(name);
    if (ifs.is_open())
    {
        ifs.close();
        return true;
    }
    else
    {
        return false;
    }
}

double xBuf[5] = {0};
double yBuf[5] = {0};

static double IIRLowPass(double x)
{
    double b[5] = {3.5884e-4, 0.001435, 0.002153, 0.001435, 3.5884e-4};
    double a[5] = {1.0, -3.2133, 3.9348, -2.1689, 0.4531};
    for (int i = 4; i > 0; i--)
    {
        yBuf[i] = yBuf[i - 1];
        xBuf[i] = xBuf[i - 1];
    }
    xBuf[0] = x;
    yBuf[0] = 0;
    for (int i = 1; i < 5; i++)
    {
        yBuf[0] = yBuf[0] + b[i] * xBuf[i];
        yBuf[0] = yBuf[0] - a[i] * yBuf[i];
    }
    yBuf[0] = yBuf[0] + b[0] * xBuf[0];
    return yBuf[0];
}

static int calZuMY(const char* file_path1, const char* file_path2, const double d5, double* M, double* Y)
{
    vector<double> tj11[4];
    vector<double> tj12[4];
    for (int i = 0; i < 4; i++)
    {
        tj11[i].reserve(3200);
        tj12[i].reserve(3200);
    }

    ifstream infile1, infile2;

    infile1.open(file_path1);
    if (!infile1.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num1 = 0;
    double data = 0;
    while (!infile1.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile1 >> data;
            tj11[i].push_back(data);
        }
        num1++;
    }
    infile1.close();

    infile2.open(file_path2);
    if (!infile2.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num2 = 0;
    while (!infile2.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile2 >> data;
            tj12[i].push_back(data);
        }
        num2++;
    }
    infile2.close();

    if (fabs(num1 - num2) > TRAJ_DIFF_NUM)
    {
        // 前后轨迹不一致
        return -1;
    }
    // 取行数更小值
    num1 = num1 < num2 ? num1 : num2;
    if (num1 < TRAJ_MULT_NUM * 2)
    {
        // 轨迹过短
        return -1;
    }

    Matrix<double, Dynamic, Dynamic> tj11fil;
    tj11fil.resize(num1, 4);
    Matrix<double, Dynamic, Dynamic> tj12fil;
    tj12fil.resize(num2, 4);
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj11fil(j, i) = IIRLowPass(tj11[i][j]); }
    }
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj12fil(j, i) = IIRLowPass(tj12[i][j]); }
    }
    double sum = 0.0;
    for (int n = 0; n < num1; n++) { sum += tj11fil(n, 3) - tj12fil(n, 3); }
    double avgm = sum / num1;
    Matrix<double, Dynamic, 1> deltm;
    deltm.resize((num1 - TRAJ_MULT_NUM * 2), 1);
    Matrix<double, Dynamic, 1> deltq;
    deltq.resize((num1 - TRAJ_MULT_NUM * 2), 1);
    for (int n = 0; n < num1 - TRAJ_MULT_NUM * 2; n++)
    {
        deltm(n, 0) = tj11fil(n + TRAJ_MULT_NUM, 1) - tj12fil(n + TRAJ_MULT_NUM, 1) + avgm;
        deltq(n, 0) = cos(tj11fil(n + TRAJ_MULT_NUM, 0)) * 9.81 * d5 / 1000.0;
    }
    *M = (deltq.transpose() * deltq).inverse() * deltq.transpose() * (deltm / MODIFY_FACTOR);
    *Y = (avgm / MODIFY_FACTOR) / *M / 9.81 * 1000;

    return 0;
}

static int calZuX(const char* file_path1, const char* file_path2, const double M, double* X)
{
    vector<double> tj21[4];
    vector<double> tj22[4];
    for (int i = 0; i < 4; i++)
    {
        tj21[i].reserve(3200);
        tj22[i].reserve(3200);
    }

    ifstream infile1, infile2;

    infile1.open(file_path1);
    if (!infile1.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num1 = 0;
    double data = 0;
    while (!infile1.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile1 >> data;
            tj21[i].push_back(data);
        }
        num1++;
    }
    infile1.close();

    infile2.open(file_path2);
    if (!infile2.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num2 = 0;
    while (!infile2.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile2 >> data;
            tj22[i].push_back(data);
        }
        num2++;
    }
    infile2.close();

    if (fabs(num1 - num2) > TRAJ_DIFF_NUM)
    {
        // 前后轨迹不一致
        return -1;
    }
    // 取行数更小值
    num1 = num1 < num2 ? num1 : num2;

    Matrix<double, Dynamic, Dynamic> tj21fil;
    tj21fil.resize(num1, 4);
    Matrix<double, Dynamic, Dynamic> tj22fil;
    tj22fil.resize(num2, 4);
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj21fil(j, i) = IIRLowPass(tj21[i][j]); }
    }
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj22fil(j, i) = IIRLowPass(tj22[i][j]); }
    }
    double sum = 0.0;
    for (int n = 0; n < num1; n++) { sum += tj21fil(n, 3) - tj22fil(n, 3); }
    double avgm = sum / num1;
    *X = (avgm / MODIFY_FACTOR) / M / 9.81 * 1000;

    return 0;
}

static int calZuZ(const char* file_path1, const char* file_path2, const double M, const double d6, double* Z)
{
    vector<double> tj31[2];
    vector<double> tj32[2];
    for (int i = 0; i < 2; i++)
    {
        tj31[i].reserve(1200);
        tj32[i].reserve(1200);
    }

    ifstream infile1, infile2;

    infile1.open(file_path1);
    if (!infile1.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num1 = 0;
    double data = 0;
    while (!infile1.eof())
    {
        for (int i = 0; i < 2; i++)
        {
            infile1 >> data;
            tj31[i].push_back(data);
        }
        num1++;
    }
    infile1.close();

    infile2.open(file_path2);
    if (!infile2.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num2 = 0;
    while (!infile2.eof())
    {
        for (int i = 0; i < 2; i++)
        {
            infile2 >> data;
            tj32[i].push_back(data);
        }
        num2++;
    }
    infile2.close();

    if (fabs(num1 - num2) > TRAJ_DIFF_NUM)
    {
        // 前后轨迹不一致
        return -1;
    }
    // 取行数更小值
    num1 = num1 < num2 ? num1 : num2;

    Matrix<double, Dynamic, Dynamic> tj31fil;
    tj31fil.resize(num1, 2);
    Matrix<double, Dynamic, Dynamic> tj32fil;
    tj32fil.resize(num2, 2);
    for (int i = 0; i < 2; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj31fil(j, i) = IIRLowPass(tj31[i][j]); }
    }
    for (int i = 0; i < 2; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj32fil(j, i) = IIRLowPass(tj32[i][j]); }
    }
    double sum = 0.0;
    for (int n = 0; n < num1; n++) { sum += tj31fil(n, 1) - tj32fil(n, 1); }
    double avgm = sum / num1;
    *Z = (avgm / MODIFY_FACTOR) / M / 9.81 * 1000.0 - d6;
    return 0;
}

static int calMcMZ(const char* file_path1, const char* file_path2, const double d4, const double d6, double* M, double* Z)
{
    vector<double> tj11[4];
    vector<double> tj12[4];
    for (int i = 0; i < 4; i++)
    {
        tj11[i].reserve(3200);
        tj12[i].reserve(3200);
    }
    ifstream infile1, infile2;

    infile1.open(file_path1);
    if (!infile1.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num1 = 0;
    double data = 0;
    while (!infile1.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile1 >> data;
            tj11[i].push_back(data);
        }
        num1++;
    }
    infile1.close();

    infile2.open(file_path2);
    if (!infile2.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num2 = 0;
    while (!infile2.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile2 >> data;
            tj12[i].push_back(data);
        }
        num2++;
    }
    infile2.close();

    if (fabs(num1 - num2) > TRAJ_DIFF_NUM)
    {
        // 前后轨迹不一致
        return -1;
    }
    // 取行数更小值
    num1 = num1 < num2 ? num1 : num2;

    if (num1 < TRAJ_MULT_NUM * 2)
    {
        // 轨迹过短
        return -1;
    }

    Matrix<double, Dynamic, Dynamic> tj11fil;
    tj11fil.resize(num1, 4);
    Matrix<double, Dynamic, Dynamic> tj12fil;
    tj12fil.resize(num2, 4);
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj11fil(j, i) = IIRLowPass(tj11[i][j]); }
    }
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj12fil(j, i) = IIRLowPass(tj12[i][j]); }
    }

    Matrix<double, Dynamic, 1> deltm;
    deltm.resize(num1 - TRAJ_MULT_NUM * 2, 1);
    Matrix<double, Dynamic, 1> deltq;
    deltq.resize(num1 - TRAJ_MULT_NUM * 2, 1);
    for (int n = 0; n < num1 - TRAJ_MULT_NUM * 2; n++)
    {
        deltm(n, 0) = -(tj11fil(n + TRAJ_MULT_NUM, 1) - tj12fil(n + TRAJ_MULT_NUM, 1)) + tj11fil(n + TRAJ_MULT_NUM, 3) - tj12fil(n + TRAJ_MULT_NUM, 3);
        deltq(n, 0) = sin(tj11fil(n + TRAJ_MULT_NUM, 0)) * 9.81 * d4 / 1000.0;
    }
    double sum = 0.0;
    for (int n = TRAJ_MULT_NUM; n < num1 - TRAJ_MULT_NUM; n++) { sum += tj11fil(n, 3) - tj12fil(n, 3); }
    double avgm = sum / (num1 - TRAJ_MULT_NUM * 2);
    *M = (deltq.transpose() * deltq).inverse() * deltq.transpose() * deltm;
    *Z = -avgm / *M / 9.81 * 1000 - d6;
    return 0;
}

static int calMcX(const char* file_path1, const char* file_path2, const double M, double* X)
{
    vector<double> tj21[4];
    vector<double> tj22[4];
    for (int i = 0; i < 4; i++)
    {
        tj21[i].reserve(3200);
        tj22[i].reserve(3200);
    }

    ifstream infile1, infile2;

    infile1.open(file_path1);
    if (!infile1.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num1 = 0;
    double data = 0;
    while (!infile1.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile1 >> data;
            tj21[i].push_back(data);
        }
        num1++;
    }
    infile1.close();

    infile2.open(file_path2);
    if (!infile2.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num2 = 0;
    while (!infile2.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile2 >> data;
            tj22[i].push_back(data);
        }
        num2++;
    }
    infile2.close();

    if (fabs(num1 - num2) > TRAJ_DIFF_NUM)
    {
        // 前后轨迹不一致
        return -1;
    }
    // 取行数更小值
    num1 = num1 < num2 ? num1 : num2;

    if (num1 < TRAJ_MULT_NUM * 2)
    {
        // 轨迹过短
        return -1;
    }

    Matrix<double, Dynamic, Dynamic> tj21fil;
    tj21fil.resize(num1, 4);
    Matrix<double, Dynamic, Dynamic> tj22fil;
    tj22fil.resize(num2, 4);
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj21fil(j, i) = IIRLowPass(tj21[i][j]); }
    }
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj22fil(j, i) = IIRLowPass(tj22[i][j]); }
    }
    double sum1 = 0.0;
    double sum2 = 0.0;
    for (int n = TRAJ_MULT_NUM; n < num1 - TRAJ_MULT_NUM; n++)
    {
        sum1 += tj21fil(n, 3) - tj22fil(n, 3);
        sum2 += tj21fil(n, 1) - tj22fil(n, 1);
    }
    double avgm = (sum1 + sum2) / 2.0 / (num1 - TRAJ_MULT_NUM * 2);
    *X = -avgm / M / 9.81 * 1000;
    return 0;
}

static int calMcY(const char* file_path1, const char* file_path2, const double M, double* Y)
{
    vector<double> tj31[4];
    vector<double> tj32[4];
    for (int i = 0; i < 4; i++)
    {
        tj31[i].reserve(3200);
        tj32[i].reserve(3200);
    }
    ifstream infile1, infile2;

    infile1.open(file_path1);
    if (!infile1.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num1 = 0;
    double data = 0;
    while (!infile1.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile1 >> data;
            tj31[i].push_back(data);
        }
        num1++;
    }
    infile1.close();

    infile2.open(file_path2);
    if (!infile2.is_open())
    {
        cout << "Error opening traj1withload.txt";
        return -1;
    }
    int num2 = 0;
    while (!infile2.eof())
    {
        for (int i = 0; i < 4; i++)
        {
            infile2 >> data;
            tj32[i].push_back(data);
        }
        num2++;
    }
    infile2.close();

    if (fabs(num1 - num2) > TRAJ_DIFF_NUM)
    {
        // 前后轨迹不一致
        return -1;
    }
    // 取行数更小值
    num1 = num1 < num2 ? num1 : num2;

    Matrix<double, Dynamic, Dynamic> tj31fil;
    tj31fil.resize(num1, 4);
    Matrix<double, Dynamic, Dynamic> tj32fil;
    tj32fil.resize(num2, 4);
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj31fil(j, i) = IIRLowPass(tj31[i][j]); }
    }
    for (int i = 0; i < 4; i++)
    {
        for (int m = 0; m < 5; m++)
        {
            xBuf[m] = 0.0;
            yBuf[m] = 0.0;
        }
        for (int j = 0; j < num1; j++) { tj32fil(j, i) = IIRLowPass(tj32[i][j]); }
    }
    double sum1 = 0.0;
    double sum2 = 0.0;
    for (int n = TRAJ_MULT_NUM; n < num1 - TRAJ_MULT_NUM; n++)
    {
        sum1 += tj31fil(n, 3) - tj32fil(n, 3);
        sum2 += tj31fil(n, 1) - tj32fil(n, 1);
    }
    double avgm = (sum1 + sum2) / 2.0 / (num1 - TRAJ_MULT_NUM * 2);
    *Y = -avgm / M / 9.81 * 1000;
    return 0;
}

int calpayload(double* dhpara, double* payload, double* center)
{
    int robotmodle = 0;
    double d4 = dhpara[0];
    double d5 = dhpara[1];
    double d6 = dhpara[2];
    // robotmodle = 1 -> cobot; robotmodle = 2 -> Zu
    if (d5 < 0.001)
    {
        robotmodle = 1;
    }
    else
    {
        robotmodle = 2;
    }
    if (fileExist("/tmp/traj1withoutload.txt") == false || fileExist("/tmp/traj2withoutload.txt") == false || fileExist("/tmp/traj3withoutload.txt") == false ||
        fileExist("/tmp/traj1withload.txt") == false || fileExist("/tmp/traj2withload.txt") == false || fileExist("/tmp/traj3withload.txt") == false)
    {
        return -1;
    }
    if (robotmodle == 2)
    {
        if (calZuMY("/tmp/traj1withload.txt", "/tmp/traj1withoutload.txt", d5, payload, &center[1]) == -1 ||
            calZuX("/tmp/traj2withload.txt", "/tmp/traj2withoutload.txt", *payload, &center[0]) == -1 ||
            calZuZ("/tmp/traj3withload.txt", "/tmp/traj3withoutload.txt", *payload, d6, &center[2]) == -1)
        {
            return -1;
        }
    }

    if (robotmodle == 1)
    {
        if (calMcMZ("/tmp/traj1withload.txt", "/tmp/traj1withoutload.txt", d4, d6, payload, &center[2]) == -1)
        {
            return -1;
        }
        *payload = 0.7536 * (*payload) + 0.0094;
        if (calMcX("/tmp/traj2withload.txt", "/tmp/traj2withoutload.txt", *payload, &center[0]) == -1 ||
            calMcY("/tmp/traj3withload.txt", "/tmp/traj3withoutload.txt", *payload, &center[1]) == -1)
        {
            return -1;
        }
    }

    if (*payload < 0.03)
    {
        *payload = 0.0;
        center[0] = 0.0;
        center[1] = 0.0;
        center[2] = 0.0;
    }

    return 0;
}